/**
 * Created by IntelliJ IDEA.
 * User: martlenn
 * Date: 25-Jan-2008
 * Time: 00:44:07
 */
package org.hkupp.util.workers;

import be.proteomics.util.sun.SwingWorker;
import be.proteomics.util.interfaces.Flamable;

import java.io.File;
import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Collection;
import java.util.ArrayList;

import org.hkupp.db.accessors.*;
import org.hkupp.gui.progressbar.DefaultProgressBar;
import org.hkupp.model.IdentifiedRun;
import org.hkupp.io.TagSummaryReader;
import org.hkupp.util.SeparatedStringSplitter;

/*
* CVS information:
*
* $Revision$
* $Date$
*/

/**
 * This class processes an array of Files for their results,
 * and stores them on the specified Connection.
 *
 * @author Lennart Martens
 * @version $Id$
 */
public class StoreDataWorker extends SwingWorker {

    private File[] iInputFiles = null;
    private Connection iConn = null;
    private Experiment iExperiment = null;
    private Instrument iInstrument = null;

    private Flamable iParent = null;
    private DefaultProgressBar iProgress = null;

    public StoreDataWorker(Connection aConn, Experiment aExperiment, File[] aInputFiles, Instrument aInstrument, Flamable aParent, DefaultProgressBar aProgress) {
        this.iConn = aConn;
        this.iExperiment = aExperiment;
        this.iInputFiles = aInputFiles;
        this.iInstrument = aInstrument;
        this.iParent = aParent;
        this.iProgress = aProgress;
    }

    /**
     * Compute the value to be returned by the <code>get</code> method.
     */
    public Object construct() {
        // The HashMap for all the proteins (shared by all runs).
        HashMap<String, Protein> proteins = new HashMap<String, Protein>();
        // Cycle through all files.
        String message = "Reading " + iInputFiles.length + " files";
        if(iProgress != null) {
            iProgress.setMessage(message + "...");
        }
        Collection<IdentifiedRun> runs = new ArrayList<IdentifiedRun>();
        for (int i = 0; i < iInputFiles.length; i++) {
            File inputFile = iInputFiles[i];
            if(iProgress != null) {
                iProgress.setMessage(message + " (" + inputFile.getParentFile().getName() + "; " + (iInputFiles.length-i) +  " files remaining)...");
            }
            try {
                runs.add(TagSummaryReader.readTagSummaryFile(inputFile, proteins));
                // Signal progress has been made.
                if(iProgress != null) {
                    iProgress.setValue(iProgress.getValue()+1);
                }
            } catch(IOException ioe) {
                ioe.printStackTrace();
            }
        }
        // Now store all data.
        storeData(runs);
        return null;
    }

    public Object get() {
        return null;
    }

    private void storeData(Collection<IdentifiedRun> aRuns) {
        // See if we have any data...
        if(aRuns == null || aRuns.size() == 0) {
            if(iProgress != null) {
                iProgress.setValue(iProgress.getMaximum());
                return;
            }
        }
        // Now we need to store all the data.
        try {
            // The Proteins are shared among all runs. So we need store these only once.
            // We'll do that here, and will keep track of their primary keys.
            // So we'll take the proteins from the first run, and work with that.
            String message = "Processing identified proteins";
            if(iProgress != null) {
                iProgress.setMessage(message + "...");
            }
            HashMap<String, Long> proteinKeys = new HashMap<String, Long>();
            IdentifiedRun run = aRuns.iterator().next();
            HashMap proteins = run.getProteins();
            int liSize = proteins.size();
            Iterator iter = proteins.keySet().iterator();
            int count = 0;
            int previousProgress= -1;
            while (iter.hasNext()) {
                String accession = (String) iter.next();
                Protein protein = Protein.findProtein(iConn, accession);
                long proteinID = -1l;
                if(protein == null) {
                    // Store it!
                    protein = (Protein)proteins.get(accession);
                    protein.persist(iConn);
                    proteinID = ((Number)protein.getGeneratedKeys()[0]).longValue();
                } else {
                    // Just get the PK.
                    proteinID = protein.getProteinid();
                }
                // Store the ID for this accession.
                proteinKeys.put(accession, proteinID);
                // See how far we've got.
                count++;
                int progress = (int)(((double)count)/liSize)*100;
                if(iProgress != null && progress > previousProgress) {
                    iProgress.setMessage(message + " (" + progress + "% complete)...");
                    previousProgress = progress;
                }
            }
            // OK, proteins done!
            if(iProgress != null) {
                iProgress.setValue(iProgress.getValue()+1);
            }

            // Other shared data is: modifications and modification types.
            HashMap<String, Long> modNameToID = new HashMap<String, Long>();
            // Load modtypes (fixed, var)
            HashMap<String, Modification_type> modTypes = Modification_type.getALlModification_types(iConn);

            message = "Processing run ";
            // Now simply cycle all runs, and store each.
            for (Iterator identifiedRunIterator = aRuns.iterator(); identifiedRunIterator.hasNext();) {
                run = (IdentifiedRun) identifiedRunIterator.next();
                String innerMessage = message + run.getLocationCode() + "-" + run.getIteration();
                if(iProgress != null) {
                    iProgress.setMessage(innerMessage + "...");
                }
                // First up: locationcode.
                long locationCodeID = -1l;
                Location_code locCode = Location_code.findLocation_code(iConn, run.getLocationCode());
                if(locCode == null) {
                    // Create a new one!
                    HashMap data = new HashMap(1);
                    data.put(Location_code.CODE, run.getLocationCode());
                    locCode = new Location_code(data);
                    locCode.persist(iConn);
                    Object[] keys = locCode.getGeneratedKeys();
                    locationCodeID = ((Number)keys[0]).longValue();
                } else {
                    locationCodeID = locCode.getLocation_codeid();
                }
                // OK, next is the location proper.
                HashMap locationData = new HashMap(3);
                locationData.put(LocationTableAccessor.L_EXPERIMENTID, iExperiment.getExperimentid());
                locationData.put(LocationTableAccessor.L_LOCATION_CODEID, locationCodeID);
                locationData.put(LocationTableAccessor.ITERATION, run.getIteration());
                LocationTableAccessor location = new LocationTableAccessor(locationData);
                location.persist(iConn);
                Object[] keys = location.getGeneratedKeys();
                long locationid = ((Number)keys[0]).longValue();

                // Right, now the peptides (who wil take care of their spectra as well).
                HashMap<File, Long> specFileToID = new HashMap<File, Long>();
                HashMap<String, Long> peptideKeys = new HashMap<String, Long>();

                HashMap peptides = run.getPeptides();
                liSize = peptides.size();
                Iterator peptideIter = peptides.keySet().iterator();
                if(iProgress != null) {
                    iProgress.setMessage(innerMessage + "0% complete...");
                }
                count = 0;
                previousProgress = -1;
                while (peptideIter.hasNext()) {
                    String peptideKey = (String)peptideIter.next();
                    Peptide peptide = (Peptide)peptides.get(peptideKey) ;
                    // See if we already have the spectrum. If not, store it first.
                    long spectrumID = -1l;
                    if(specFileToID.containsKey(peptide.getSpectrumFile())) {
                        spectrumID = specFileToID.get(peptide.getSpectrumFile()).longValue();
                    } else {
                        // Create a new spectrum, and store.
                        HashMap specData = new HashMap(4);
                        specData.put(Spectrum.L_LOCATIONID, locationid);
                        specData.put(Spectrum.L_INSTRUMENTID, iInstrument.getInstrumentid());
                        Spectrum spectrum = new Spectrum(specData);
                        try {
                            spectrum.loadSpectrumFile(peptide.getSpectrumFile());
                        } catch(IOException ioe) {
                            System.err.println("Unable to load spectrum '" + peptide.getSpectrumFile() + "'!");
                            ioe.printStackTrace();
                        }
                        spectrum.persist(iConn);
                        Object[] specKeys = spectrum.getGeneratedKeys();
                        spectrumID = ((Number)specKeys[0]).longValue();
                        // Add spectrumID to lookup hashmap.
                        specFileToID.put(peptide.getSpectrumFile(), spectrumID);
                    }
                    // OK, spectrum reference found.
                    peptide.setL_spectrumid(spectrumID);
                    // Swap sequence to modified sequence at this point.
                    peptide.setModified_sequence(peptide.getSequence());
                    peptide.setSequence(peptide.getSequence().toUpperCase());
                    // Store peptide.
                    peptide.persist(iConn);
                    long peptideID = ((Number)peptide.getGeneratedKeys()[0]).longValue();
                    peptideKeys.put(peptideKey, peptideID);
                    // Now check the modifications.
                    // First the variable mods.
                    String varMods = peptide.getVarMods();
                    long modificationTypeID = modTypes.get("variable modification").getModification_typeid();
                    String modSeq = peptide.getModified_sequence();
                    if(varMods != null && !varMods.trim().equals("")) {
                        SeparatedStringSplitter ssp = new SeparatedStringSplitter(varMods.trim(), " ", true);
                        String[] varModifications = ssp.allTokens();
                        for (int i = 0; i < varModifications.length; i++) {
                            String varModification = varModifications[i];
                            String code = varModification.substring(0, varModification.indexOf(":")).trim();
                            String name = varModification.substring(varModification.indexOf(":")+1).trim();
                            long modificationID = -1;
                            if(modNameToID.containsKey(name)) {
                                modificationID = modNameToID.get(name);
                            } else {
                                // Check whether it is in the DB.
                                Modification mod = Modification.findModification(iConn, name);
                                if(mod != null) {
                                    // Found it!
                                    modificationID = mod.getModificationid();
                                } else {
                                    // Create a new one.
                                    HashMap modData = new HashMap(1);
                                    modData.put(Modification.NAME, name);
                                    mod = new Modification(modData);
                                    mod.persist(iConn);
                                    modificationID = ((Long)mod.getGeneratedKeys()[0]).longValue();
                                }
                                // Add the now known mod to the HashMap.
                                modNameToID.put(name, modificationID);
                            }
                            // Find all locations for this modification.
                            int prevLoc = 0;
                            int modLoc = -1;
                            while((modLoc = modSeq.indexOf(code, prevLoc)) >= 0) {
                                HashMap modPepData = new HashMap(4);
                                modPepData.put(Modification_to_peptideTableAccessor.L_MODIFICATION_TYPEID, modificationTypeID);
                                modPepData.put(Modification_to_peptideTableAccessor.L_MODIFICATIONID, modificationID);
                                modPepData.put(Modification_to_peptideTableAccessor.L_PEPTIDEID, peptideID);
                                modPepData.put(Modification_to_peptideTableAccessor.LOCATION, new Long(modLoc));
                                Modification_to_peptideTableAccessor mpta = new Modification_to_peptideTableAccessor(modPepData);
                                mpta.persist(iConn);
                                prevLoc = modLoc+1;
                            }
                        }
                    }

                    // Now the fixed modification.
                    String fixedMods = peptide.getFixedMods();
                    modificationTypeID = modTypes.get("fixed modification").getModification_typeid();
                    modSeq = peptide.getModified_sequence();
                    if(fixedMods != null && !fixedMods.trim().equals("")) {
                        SeparatedStringSplitter ssp = new SeparatedStringSplitter(fixedMods.trim(), " ", true);
                        String[] fixedModifications = ssp.allTokens();
                        for (int i = 0; i < fixedModifications.length; i++) {
                            String fixedModification = fixedModifications[i];
                            String code = fixedModification.substring(0, fixedModification.indexOf(":")).trim();
                            String name = fixedModification.substring(fixedModification.indexOf(":")+1).trim();
                            long modificationID = -1;
                            if(modNameToID.containsKey(name)) {
                                modificationID = modNameToID.get(name);
                            } else {
                                // Check whether it is in the DB.
                                Modification mod = Modification.findModification(iConn, name);
                                if(mod != null) {
                                    // Found it!
                                    modificationID = mod.getModificationid();
                                } else {
                                    // Create a new one.
                                    HashMap modData = new HashMap(1);
                                    modData.put(Modification.NAME, name);
                                    mod = new Modification(modData);
                                    mod.persist(iConn);
                                    modificationID = ((Long)mod.getGeneratedKeys()[0]).longValue();
                                }
                                // Add the now known mod to the HashMap.
                                modNameToID.put(name, modificationID);
                            }
                            // Find all locations for this modification.
                            int prevLoc = 0;
                            int modLoc = -1;
                            while((modLoc = modSeq.indexOf(code, prevLoc)) >= 0) {
                                HashMap modPepData = new HashMap(4);
                                modPepData.put(Modification_to_peptideTableAccessor.L_MODIFICATION_TYPEID, modificationTypeID);
                                modPepData.put(Modification_to_peptideTableAccessor.L_MODIFICATIONID, modificationID);
                                modPepData.put(Modification_to_peptideTableAccessor.L_PEPTIDEID, peptideID);
                                modPepData.put(Modification_to_peptideTableAccessor.LOCATION, new Long(modLoc));
                                Modification_to_peptideTableAccessor mpta = new Modification_to_peptideTableAccessor(modPepData);
                                mpta.persist(iConn);
                                prevLoc = modLoc+1;
                            }
                        }
                    }
                    // See how far we've got.
                    count++;
                    int progress = (int)(((double)count)/liSize)*100;
                    if(iProgress != null && progress > previousProgress) {
                        iProgress.setMessage(innerMessage + " (" + progress + "% complete)...");
                        previousProgress = progress;
                    }
                }

                // Finally, the peptide to protein mapping.
                HashMap peptidesToProteins = run.getPeptideToProteinMap();
                HashMap pepLocations = run.getPeptideLocationInProtein();
                iter = peptidesToProteins.keySet().iterator();
                if(iProgress != null) {
                    iProgress.setMessage(innerMessage + " (peptide to protein mapping)...");
                }
                while (iter.hasNext()) {
                    String peptideKey = (String) iter.next();
                    String sequence = peptideKey.substring(peptideKey.indexOf("@")+1).trim();
                    Collection matchedProteins = (Collection)peptidesToProteins.get(peptideKey);
                    for (Iterator iterator = matchedProteins.iterator(); iterator.hasNext();) {
                        String accession = (String) iterator.next();
                        // Get the protein ID.
                        long proteinID = proteinKeys.get(accession);
                        // Get the peptide ID.
                        long peptideID = peptideKeys.get(peptideKey);

                        // Cycle all known locations for this protein.
                        Iterator<Integer> locIter = ((Collection<Integer>)pepLocations.get(sequence + "@" + accession)).iterator();
                        while (locIter.hasNext()) {
                            Integer startLocation = locIter.next();
                            // Calculate stop position.
                            Integer stopLocation = startLocation + sequence.length();

                            HashMap pepToProtData = new HashMap(4);
                            pepToProtData.put(Peptide_to_proteinTableAccessor.L_PEPTIDEID, peptideID);
                            pepToProtData.put(Peptide_to_proteinTableAccessor.L_PROTEINID, proteinID);
                            pepToProtData.put(Peptide_to_proteinTableAccessor.START_POSITION, startLocation.longValue());
                            pepToProtData.put(Peptide_to_proteinTableAccessor.END_POSITION, stopLocation.longValue());
                            Peptide_to_proteinTableAccessor ptpta = new Peptide_to_proteinTableAccessor(pepToProtData);
                            ptpta.persist(iConn);
                        }
                    }
                }
                if(iProgress != null) {
                    iProgress.setValue(iProgress.getValue()+1);
                }
            }
        } catch(SQLException sqle) {
            sqle.printStackTrace();
        }
        // All done!
        if(iProgress != null) {
            iProgress.setValue(iProgress.getMaximum());
        }
    }

}
